# !/usr/bin/env python3
# -*- coding: utf-8 -*-
# @Time : 2023/4/24 11:25
# @Author : xyjiang1216
# @Email : jiangxuyan@nudt.edu.cn
# @File : layer_4_preprocessing_input.py
# @Version: OpenPlanner2.0
import importlib
import json
import math

from OpenPlanner.core.api.algo_input_class import StreamInfo
from OpenPlanner.core.src.op_class import OPClass
from OpenPlanner.debug.debug_utils import DEBUG_PRINT
from OpenPlanner.util.compute_hyper_period import compute_hyper_period


class InputProcessingModule:
    def __init__(self,
                 op_obj: OPClass):
        self.op_obj = op_obj
        self.input_hook_func = None

    def _init_link_array(self):
        # 从topo_file中读取link_set
        topo_filepath = self.op_obj.topo_filepath
        with open(topo_filepath, 'r') as fd:
            link_set = json.load(fd)
        # 判断用户是否是自定义Link类
        user_defined_link = self.op_obj.init_input.user_defined_link
        link_name = 'Link'
        LinkClass = getattr(
            importlib.import_module('OpenPlanner.core.api.algo_input_class'),
            link_name
        )
        if user_defined_link != '':
            # 动态加载自定义链路
            class_name = user_defined_link
            base_module = self.op_obj.init_input.base_module_path
            try:
                module = importlib.import_module(base_module)
            except ModuleNotFoundError:
                print('Error: 用户自定义模块%s不存在，'
                      '请检查路径后更新配置文件' % base_module)
                exit(0)
            try:
                LinkClass = getattr(module, class_name)
            except AttributeError:
                print('AttributeError: 用户自定义模块%s'
                      '没有%s属性，'
                      '请检查路径后更新配置文件'
                      % (base_module, class_name))
                exit(0)

        link_array = []
        for link in link_set:
            link_obj = LinkClass(**link)
            link_array.append(link_obj)

        self.op_obj.algo_input.link_array = link_array

    def _init_stream_array(self):
        stream_filepath = self.op_obj.stream_filepath
        with open(stream_filepath, 'r') as fd:
            stream_set = json.load(fd)
        # 判断用户是否自定义Stream类
        user_defined_stream = \
            self.op_obj.init_input.user_defined_stream
        StreamClass = getattr(
            importlib.import_module('OpenPlanner.core.api.algo_input_class'),
            'Stream'
        )
        if user_defined_stream != '':
            # 动态加载自定义流量类
            class_name = self.op_obj.init_input.user_defined_stream
            base_module = self.op_obj.init_input.base_module_path
            try:
                module = importlib.import_module(base_module)
            except ModuleNotFoundError:
                print('ModuleNotFoundError: 用户自定义流量模块'
                      '%s不存在，请检查模块路径后'
                      '更新配置文件' % base_module)
                exit(0)

            try:
                StreamClass = getattr(module, class_name)
            except AttributeError:
                print('AttributeError: 用户自定义模块%s'
                      '中不存在%s类，请检查路径后更新配置文件'
                      % (base_module, class_name))
                exit(0)

        stream_array = []
        for stream in stream_set:
            stream_obj = StreamClass(**stream)
            stream_array.append(stream_obj)
            hop_id = 0
            for link_id in stream['route']:
                link_obj = \
                    self.op_obj.algo_input.link_array[link_id]
                link_obj.add_stream_to_link(
                    StreamInfo(
                        **{"stream_id": stream['stream_id'],
                           "hop_id": hop_id}
                    )
                )
                hop_id += 1
        self.op_obj.algo_input.stream_array = stream_array

    def _init_instance_array(self):
        user_defined_instance = \
            self.op_obj.init_input.user_defined_instance
        ClassName = getattr(
            importlib.import_module('OpenPlanner.core.api.algo_input_class'),
            'Instance'
        )
        if user_defined_instance != '':
            # 动态加载自定义流实例类
            class_name = user_defined_instance
            base_module = \
                self.op_obj.init_input.base_module_path
            try:
                module = importlib.import_module(
                    base_module
                )
            except ModuleNotFoundError:
                print("ModuleNotFoundError: 用户自定义模块%s不存在，"
                      "请检查路径后更新配置文件" % base_module)
                exit(0)
            try:
                ClassName = getattr(module, class_name)
            except AttributeError:
                print("AttributeError: 模块%s中不存在"
                      "%s类，请检查路径后更新配置文件"
                      % (base_module, class_name))
                exit(0)

        hyper_period = compute_hyper_period(
            *[stream_obj.period for stream_obj in \
              self.op_obj.algo_input.stream_array]
        )

        stream_array = self.op_obj.algo_input.stream_array
        instance_obj_set_per_stream = []
        for stream_obj in stream_array:
            period = stream_obj.period
            route_set = stream_obj.route
            stream_id = stream_obj.stream_id

            instance_obj_set_per_link = []

            for link_id in route_set:
                instance_obj_per_period = []
                for period_id in \
                        range(math.ceil(hyper_period / period)):
                    kargs = {"stream_id": stream_id,
                             "period_id": period_id,
                             "link_id": link_id}

                    stream_instance_obj = ClassName(**kargs)
                    instance_obj_per_period.append(stream_instance_obj)
                instance_obj_set_per_link.append(instance_obj_per_period)
            instance_obj_set_per_stream.append(instance_obj_set_per_link)
        self.op_obj.algo_input.instance_array = instance_obj_set_per_stream

    def init_input(self):
        self._init_link_array()
        self._init_stream_array()
        self._init_instance_array()

    # 钩子函数相关
    def _register_input_hook(self, func):
        self.input_hook_func = func

    def custom_input_hook(self):
        base_module = self.op_obj.input_init_hook.base_module_path
        func_name = self.op_obj.input_init_hook.hook_func_name

        if base_module != '' and func_name != '':
            # 注册钩子
            # 判断模块是否存在
            try:
                module = importlib.import_module(base_module)
            except ModuleNotFoundError:
                print("ModuleNotFoundError:"
                      "输入预处理钩子函数的%s模块不存在，"
                      "请检查钩子函数路径后更新配置文件" % base_module)
                exit(0)
            try:
                func = getattr(module, func_name)
            except AttributeError:
                print("AttributeError:"
                      "模块%s内没有函数%s，"
                      "请检查输入预处理钩子函数名称后更新配置文件"
                      % (base_module, func_name))
                exit(0)
            self._register_input_hook(func)

        # 执行钩子函数
        if self.input_hook_func is not None:
            kwargs = self.op_obj.input_init_hook.hook_func_param
            self.input_hook_func(self.op_obj, **kwargs)
        else:
            DEBUG_PRINT("No hook after pre-processing input")
            pass